#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional
import asyncio
import numpy as np
import aiohttp

from app.core.logging import get_logger
from app.core.exceptions import EmbeddingError
from app.core.config import settings

# 创建日志记录器
logger = get_logger("embedding")

# 创建一个简单的嵌入管理器
class SimpleEmbeddingManager:
    """简单的嵌入管理器
    
    用于替代原有的嵌入管理器
    """
    
    def __init__(self):
        """初始化嵌入管理器"""
        self.dimension = settings.EMBEDDING_DIMENSION
        self.ollama_base_url = settings.OLLAMA_BASE_URL
        self.ollama_model = settings.OLLAMA_EMBEDDING_MODEL
        
        # 确保 ollama_base_url 不为空
        if not self.ollama_base_url:
            logger.warning(f"Ollama 服务 URL 未设置，使用默认值: http://localhost:11434")
            self.ollama_base_url = "http://localhost:11434"
        
        logger.info(f"已初始化简单嵌入管理器，维度: {self.dimension}")
        logger.info(f"Ollama 服务 URL: {self.ollama_base_url}")
        logger.info(f"Ollama 嵌入模型: {self.ollama_model}")
    
    async def embed_text(self, text: str) -> List[float]:
        """对文本进行向量化"""
        if not text.strip():
            logger.warning("输入文本为空，返回零向量")
            return [0.0] * self.dimension
        
        try:
            # 确保 ollama_base_url 不为空
            if not self.ollama_base_url:
                raise EmbeddingError("Ollama 服务 URL 未设置")
            
            logger.info(f"使用 Ollama 服务: {self.ollama_base_url}")
            url = f"{self.ollama_base_url}/api/embed"
            payload = {
                "model": self.ollama_model,
                "input": [text]
            }
            
            logger.debug(f"Ollama 嵌入请求: {url}, 模型: {self.ollama_model}")
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as response:
                    if response.status != 200:
                        error_msg = f"Ollama向量化请求失败: {response.status}"
                        logger.error(error_msg)
                        raise EmbeddingError(error_msg)
                    
                    result = await response.json()
                    
                    if not result.get('embeddings'):
                        raise EmbeddingError("Ollama向量化响应格式错误")
                    
                    embeddings = result['embeddings']
                    if not embeddings:
                        raise EmbeddingError("未获取到向量数据")
                    
                    # 获取第一个文本的向量
                    vector = embeddings[0]
                    
                    logger.debug(f"Ollama文本向量化成功，维度: {len(vector)}")
                    return vector
                    
        except Exception as e:
            error_msg = f"文本向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """对多个文本进行批量向量化"""
        if not texts:
            return []
        
        # 过滤空文本
        non_empty_texts = [text for text in texts if text.strip()]
        if not non_empty_texts:
            logger.warning("所有输入文本都为空，返回零向量列表")
            return [[0.0] * self.dimension] * len(texts)
        
        try:
            # 确保 ollama_base_url 不为空
            if not self.ollama_base_url:
                raise EmbeddingError("Ollama 服务 URL 未设置")
            
            logger.info(f"使用 Ollama 服务: {self.ollama_base_url}")
            # Ollama可以一次处理多个文本
            url = f"{self.ollama_base_url}/api/embed"
            payload = {
                "model": self.ollama_model,
                "input": non_empty_texts
            }
            
            logger.debug(f"Ollama 批量嵌入请求: {url}, 模型: {self.ollama_model}, 文本数量: {len(non_empty_texts)}")
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as response:
                    if response.status != 200:
                        error_msg = f"Ollama批量向量化请求失败: {response.status}"
                        logger.error(error_msg)
                        raise EmbeddingError(error_msg)
                    
                    result = await response.json()
                    
                    if not result.get('embeddings'):
                        raise EmbeddingError("Ollama批量向量化响应格式错误")
                    
                    all_vectors = result['embeddings']
                    
                    # 处理原始文本中的空文本
                    result_vectors = []
                    vector_index = 0
                    
                    for text in texts:
                        if text.strip():
                            result_vectors.append(all_vectors[vector_index])
                            vector_index += 1
                        else:
                            result_vectors.append([0.0] * self.dimension)
                    
                    logger.info(f"Ollama批量向量化成功，处理了{len(texts)}个文本")
                    return result_vectors
                    
        except Exception as e:
            error_msg = f"批量向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_query(self, query: str) -> List[float]:
        """对查询文本进行向量化"""
        processed_query = f"查询: {query}" if len(query) < 100 else query
        return await self.embed_text(processed_query)
    
    def calculate_similarity(self, vector1: List[float], vector2: List[float]) -> float:
        """计算两个向量的余弦相似度"""
        try:
            v1 = np.array(vector1)
            v2 = np.array(vector2)
            
            # 计算余弦相似度
            dot_product = np.dot(v1, v2)
            norm1 = np.linalg.norm(v1)
            norm2 = np.linalg.norm(v2)
            
            if norm1 == 0 or norm2 == 0:
                return 0.0
            
            similarity = dot_product / (norm1 * norm2)
            return float(similarity)
            
        except Exception as e:
            logger.error(f"计算相似度失败: {str(e)}")
            return 0.0

# 创建嵌入管理器实例
base_embedding_manager = SimpleEmbeddingManager()


class EmbeddingService:
    """嵌入服务
    
    封装现有的嵌入管理器，提供额外的功能
    """
    
    def __init__(self):
        """初始化嵌入服务"""
        self.embedding_manager = base_embedding_manager
        logger.info("已初始化嵌入服务")
    
    async def embed_text(self, text: str) -> List[float]:
        """对文本进行向量化
        
        Args:
            text: 要向量化的文本
            
        Returns:
            List[float]: 文本的向量表示
            
        Raises:
            EmbeddingError: 当向量化失败时
        """
        try:
            return await self.embedding_manager.embed_text(text)
        except Exception as e:
            error_msg = f"文本向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """对多个文本进行批量向量化
        
        Args:
            texts: 要向量化的文本列表
            
        Returns:
            List[List[float]]: 文本向量列表
            
        Raises:
            EmbeddingError: 当向量化失败时
        """
        try:
            return await self.embedding_manager.embed_texts(texts)
        except Exception as e:
            error_msg = f"批量文本向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def embed_query(self, query: str) -> List[float]:
        """对查询文本进行向量化
        
        Args:
            query: 查询文本
            
        Returns:
            List[float]: 查询向量
            
        Raises:
            EmbeddingError: 当向量化失败时
        """
        try:
            return await self.embedding_manager.embed_query(query)
        except Exception as e:
            error_msg = f"查询向量化失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    def calculate_similarity(self, vector1: List[float], vector2: List[float]) -> float:
        """计算两个向量的余弦相似度
        
        Args:
            vector1: 第一个向量
            vector2: 第二个向量
            
        Returns:
            float: 余弦相似度
            
        Raises:
            EmbeddingError: 当计算失败时
        """
        try:
            return self.embedding_manager.calculate_similarity(vector1, vector2)
        except Exception as e:
            error_msg = f"计算相似度失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        try:
            return self.embedding_manager.get_model_info()
        except Exception as e:
            error_msg = f"获取模型信息失败: {str(e)}"
            logger.error(error_msg)
            raise EmbeddingError(error_msg)
    
    async def test_connection(self) -> bool:
        """测试向量化服务连接
        
        Returns:
            bool: 连接是否正常
        """
        try:
            return await self.embedding_manager.test_connection()
        except Exception as e:
            error_msg = f"测试连接失败: {str(e)}"
            logger.error(error_msg)
            return False


# 全局嵌入服务实例
embedding_service = EmbeddingService()


async def get_embedding_manager() -> EmbeddingService:
    """获取嵌入服务
    
    Returns:
        EmbeddingService: 嵌入服务实例
    """
    return embedding_service